package xconn

import (
	"bufio"
	"context"
	"encoding/binary"
	"io"
	"net"
	"sieve_admin_server_userpass/lib/app"
	"sieve_admin_server_userpass/lib/xconn/packet"
	"sync"
	"time"
)

type Conn struct {
	conn    net.Conn
	rbuf    io.Reader
	pool    *reqPool
	handler func([]byte) []byte
	cfg     *Config
	once    sync.Once
}

const (
	typeMsg  byte = '*'
	typeReq  byte = '>'
	typeResp byte = '<'
)

type RawRW interface {
	Read() ([]byte, error)
	Write(data []byte) error
}

type Config struct {
	OnInit   func(rw RawRW) error
	OnClosed func()
	Timeout  time.Duration
}

var DefaultTimeout = time.Second * 3

func NewConn(conn net.Conn, cfg *Config) (*Conn, error) {
	c := &Conn{
		conn: conn,
		rbuf: bufio.NewReader(conn),
		pool: newPool(),
		cfg:  cfg,
	}
	if c.cfg == nil {
		c.cfg = &Config{}
	}
	if c.cfg.Timeout == 0 {
		c.cfg.Timeout = DefaultTimeout
	}
	if c.cfg.OnInit != nil {
		if err := c.cfg.OnInit(&rawRW{c}); err != nil {
			c.Close()
			return nil, err
		}
	}
	app.Async(c.recvLoop)
	return c, nil
}

func (c *Conn) Close() { c.once.Do(c.close) }
func (c *Conn) close() {
	c.conn.Close()
	if c.cfg.OnClosed != nil {
		c.cfg.OnClosed()
	}
}

func (c *Conn) Handle(handler func([]byte) []byte) {
	c.handler = handler
}

func (c *Conn) send(data []byte) error {
	return packet.WritePacket(c.conn, data)
}

func (c *Conn) recvLoop() {
	defer c.Close()
	for {
		msg, err := packet.ReadPacket(c.rbuf)
		if err != nil || len(msg) == 0 {
			break
		}
		func(msg []byte) {
			app.Async(func() {
				switch msg[0] {
				case typeReq:
					if len(msg) < 5 {
						return
					}
					if c.handler != nil {
						if ret := c.handler(msg[5:]); ret != nil {
							resp := make([]byte, 5+len(ret))
							resp[0] = typeResp
							copy(resp[1:5], msg[1:5])
							copy(resp[5:], ret)
							c.send(resp)
						}
					}

				case typeResp:
					if len(msg) < 5 {
						return
					}
					c.pool.resp(binary.BigEndian.Uint32(msg[1:5]), msg[5:])

				case typeMsg:
					if c.handler != nil {
						c.handler(msg[1:])
					}
				}
			})
		}(msg)
	}
}

func (c *Conn) Post(data []byte) error {
	b := make([]byte, 1+len(data))
	b[0] = typeMsg
	copy(b[1:], data)
	return c.send(b)
}

func (c *Conn) Request(data []byte) ([]byte, error) {
	return c.RequestTimeout(c.cfg.Timeout, data)
}

func (c *Conn) RequestTimeout(timeout time.Duration, data []byte) ([]byte, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	return c.RequestCtx(ctx, data)
}

func (c *Conn) RequestCtx(ctx context.Context, data []byte) ([]byte, error) {
	seq, ch := c.pool.add()

	b := make([]byte, 5+len(data))
	b[0] = typeReq
	binary.BigEndian.PutUint32(b[1:5], seq)
	copy(b[5:], data)
	if err := c.send(b); err != nil {
		return nil, err
	}
	select {
	case <-ctx.Done():
		c.pool.del(seq)
		return nil, ctx.Err()
	case ret := <-ch:
		return ret, nil
	}
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

//

type reqPool struct {
	seq  uint32
	lock sync.Mutex
	reqs map[uint32]chan []byte
}

func newPool() *reqPool {
	return &reqPool{reqs: map[uint32]chan []byte{}}
}

const maxUint32 = uint32(1<<32 - 1)

func (r *reqPool) add() (uint32, chan []byte) {
	r.lock.Lock()
	defer r.lock.Unlock()

	r.seq += 1
	if r.seq == maxUint32 {
		r.seq = 1
	}
	ch := make(chan []byte)
	r.reqs[r.seq] = ch

	return r.seq, ch
}

func (r *reqPool) del(seq uint32) chan []byte {
	r.lock.Lock()
	defer r.lock.Unlock()

	if ch, ok := r.reqs[seq]; ok {
		delete(r.reqs, seq)
		return ch
	}
	return nil
}

func (r *reqPool) resp(seq uint32, data []byte) {
	ch := r.del(seq)
	if ch != nil {
		ch <- data
	}
}

//

type rawRW struct {
	c *Conn
}

func (r *rawRW) Read() ([]byte, error) {
	return packet.ReadPacket(r.c.rbuf)
}

func (r *rawRW) Write(data []byte) error {
	return packet.WritePacket(r.c.conn, data)
}
